import os
import cv2
import mycv
import time

import RotateImg
import utils
import random
# import RotateImg
import numpy as np
from tqdm import tqdm


def contrast_img(img_src, c, b):  # 亮度就是每个像素所有通道都加上b

    # 新建全零(黑色)图片数组:np.zeros(img1.shape, dtype=uint8)
    blank = np.zeros(img_src.shape, img_src.dtype)
    dst = cv2.addWeighted(img_src, c, blank, 1 - c, b)
    return dst


def get_label_size():
    r = random.random()
    if r < 0.3:
        return random.randint(30, 35)
    elif random.random() < 0.8:
        return random.randint(30, 50)
    else:
        return random.randint(50, 70)


# ======================================================================函数功能
# 将图片按MinMultiple到 MaxMultiple倍label的大小随机裁剪

# ======================================================================参数设置begin
# 若要随机的宽高(MinMultiple, MaxMultiple)倍，将WH设为0即可

SrcSize = (1200, 450)
DstSize = (900, 300)
MinMultiple = 8
MaxMultiple = 12

SrcDir = r'/media/fang/TOSHIBA EXT/Temp/Circle/RemoveLimit/SmallLabel'  # 源文件路径
DstDir = r'/media/fang/TOSHIBA EXT/Temp/900x300/R'  # 输出文件路径
BackGroundPicPath = r'/media/fang/TOSHIBA EXT/BackGroundpic/road'
Resize = False
GrayPadding = True
Rotate = True
Blur = False
Contrast = False
ResizeResult = True
red2black = False
black2red = False
ReviseSize = False
RotateRange = [-5, 5]
Epochs = 3
NameNum = 0


names = []
OverLargeList = []
OverLarge = 0

# 删除目的文件夹原有内容
# for files in os.scandir(DstDir):
#     os.remove(DstDir + "\\" + files.name)

# for file in os.scandir(SrcDir):
#     if file.name[-1] == 'l':
#         names.append(file.name)

BackGroundPics = []
for root, dirs, files in os.walk(BackGroundPicPath):
    for file in files:
        BackGroundPics.append(root + '/' + file)
for root, dirs, files in os.walk(SrcDir):
    for file in files:
        if file[-1] == 'l':
            names.append(root + "/" + file)

LenNames = len(names)
for Epoch in range(Epochs):
    for FileIdx, file_name in enumerate(tqdm(names)):
        WHScale = random.uniform(0.9, 1.2)
        W = int(SrcSize[0] * WHScale)
        H = int(SrcSize[1] * WHScale)
        img = cv2.imread(file_name[:-4] + '.jpg')
        # img_cp = img.copy()

        locations = utils.ReadXml(file_name)

        if Resize:
            img, locations = utils.ResizeJpgAndXml(img, locations, 1920, 1080)

        if Rotate:
            img, locations = RotateImg.RotateImgXml(img, locations, RotateRange)

        i_h, i_w, _ = img.shape
        for Info in locations:
            SizeRateW = random.uniform(MinMultiple, MaxMultiple)
            SizeRateH = random.uniform(MinMultiple, MaxMultiple)

            if W != 0 or H != 0:
                if ReviseSize:
                    label_size = get_label_size()
                    rate_size = min(Info[2] - Info[0], Info[3] - Info[1]) / label_size
                    if rate_size > 1:
                        rate_size = max(rate_size, 1.3)
                    else:
                        rate_size = min(rate_size, 0.8)
                    CutPicW = W * rate_size
                    CutPicH = H * rate_size
                else:
                    CutPicW = W
                    CutPicH = H
            else:
                    CutPicW = int((Info[2] - Info[0]) * SizeRateW)
                    CutPicH = int((Info[3] - Info[1]) * SizeRateH)

            CutPicW = int(CutPicW * random.uniform(0.95, 1.05))
            CutPicH = int(CutPicH * random.uniform(0.95, 1.05))
            # Overlarge
            if GrayPadding is False:
                if CutPicW > i_w or CutPicH > i_h or CutPicW < (
                        Info[2] - Info[0]) + 1 or CutPicH < (Info[3] - Info[1]) + 1:
                    OverLarge += 1
                    if file_name not in OverLargeList:
                        OverLargeList.append(file_name)
                    continue  # 标签图片过大

            if max(Info[2] - Info[0], Info[3] - Info[1]) > 250:
                OverLarge += 1
                if file_name not in OverLargeList:
                    OverLargeList.append(file_name)
                continue
            try:
                RandW = random.randint(0, CutPicW - (Info[2] - Info[0]))
                RandH = random.randint(0, CutPicH - (Info[3] - Info[1]))
            except ValueError:
                print("\n------------------------------------------------RandError"
                      "-----------------------------------------")
                print("The file name is %s" % file_name)
                continue
            NameNum += 1
            NewLocation0 = DstDir + "/T%05d.xml" % NameNum
            if GrayPadding:
                # ImgGrayStripe = np.ones((i_h + CutPicH * 2, i_w + CutPicW * 2, 3), dtype=np.uint8) * 128
                ImgGrayStripe = cv2.resize(cv2.imread(random.choice(BackGroundPics)), (i_w + CutPicW * 2, i_h + CutPicH * 2))
                ImgGrayStripe[CutPicH:CutPicH + i_h, CutPicW:CutPicW + i_w] = img
                imgNew = ImgGrayStripe[Info[1] - RandH + CutPicH:Info[1] - RandH + CutPicH * 2,
                         Info[0] - RandW + CutPicW:Info[0] - RandW + CutPicW * 2]

            else:
                while RandW > Info[0] or RandH > Info[1] or RandW + i_w - Info[0] < CutPicW or RandH + i_h - \
                        Info[1] < CutPicH:
                    if RandW > Info[0]:
                        RandW = random.randint(0, Info[0])
                    if RandH > Info[1]:
                        RandH = random.randint(0, Info[1])
                    if RandW + i_w - Info[0] < CutPicW:
                        RandW = random.randint(CutPicW - i_w + Info[0], CutPicW - (Info[2] - Info[0]))
                    if RandH + i_h - Info[1] < CutPicH:
                        RandH = random.randint(CutPicH - i_h + Info[1], CutPicH - (Info[3] - Info[1]))

                imgNew = img[Info[1] - RandH:Info[1] - RandH + CutPicH,
                         Info[0] - RandW:Info[0] - RandW + CutPicW]

            NewLocations = utils.CutInfos(Info[0] - RandW, Info[1] - RandH, Info[0] - RandW + CutPicW,
                                          Info[1] - RandH + CutPicH, NewLocation0, locations, CoverRate=0.15)

            if red2black:
                for location in NewLocations[1:]:
                    if random.random() > 0.7 and location[3] - location[1] > 40:
                        color = mycv.get_black_pixel()
                        imgNew[location[1]: location[3], location[0]: location[2]] = mycv.red2black_circle(
                            imgNew[location[1]: location[3], location[0]: location[2]], color)
            if black2red:
                for location in NewLocations[1:]:
                    if random.random() > 0.7 and location[3] - location[1] > 40:
                        color = mycv.get_red_pixel2()
                        imgNew[location[1]: location[3], location[0]: location[2]] = mycv.black2red_circle(
                            imgNew[location[1]: location[3], location[0]: location[2]], color)

            if Blur:
                noise = np.random.randint(0, 256, size=imgNew.shape, dtype=np.uint8)
                noise = cv2.GaussianBlur(noise, (5, 5), 0)
                imgNew = cv2.addWeighted(imgNew, 0.9, noise, 0.1, 0)
                # imgNew = cv2.GaussianBlur(imgNew, (5, 5), 0)
                # imgNew = cv2.medianBlur(imgNew, 5)
                # imgNew = cv2.medianBlur(imgNew, 5)

            if ResizeResult:
                # if imgNew.shape[0] != H:
                imgNew, NewLocations[1:] = utils.ResizeJpgAndXml(imgNew, NewLocations[1:], DstSize[0], DstSize[1])

            if Contrast:
                imgNew = contrast_img(imgNew, random.uniform(0.8, 1.8), random.randint(-10, 10))

            cv2.imwrite(NewLocation0[:-3] + "jpg", imgNew)
            utils.WriteXml(NewLocations, imgNew.shape[1], imgNew.shape[0])

    print("-----------------------输出%d张图片----------------------------" % NameNum)
    print("标签过大而放弃的图片 %d 张" % OverLarge)
    time.sleep(0.1)

# RenameJpgAndXml(DstDir)
print("------------------------------done-------------------------------")
